{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bda9eda8b4566a0d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Runnable as Action\n",
    "\n",
    "This guide will teach you how to use a `Runnable` as an action inside a guardrails configuration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a5ddc8b17af62afa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:11.284164Z",
     "start_time": "2024-01-25T14:27:11.025161Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Init: remove any existing configuration\n",
    "!rm -r config\n",
    "!mkdir config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "724db36201c3d409",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Prerequisites\n",
    "\n",
    "Set up an OpenAI API key, if not already set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4e52b23b90077cf4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:11.418023Z",
     "start_time": "2024-01-25T14:27:11.286549Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!export OPENAI_API_KEY=$OPENAI_API_KEY    # Replace with your own key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e562d3428d331b96",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Install the LangChain x OpenAI integration package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a335303d80b3953",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!pip install langchain-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b6fb59034bcb2bb",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "If you're running this inside a notebook, you also need to patch the AsyncIO loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7ba19d5c8bdc57a3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:13.693091Z",
     "start_time": "2024-01-25T14:27:13.686555Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b27d3fa09bbe91",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Sample Runnable\n",
    "\n",
    "Let's create a sample `Runnable` that checks if a string provided as input contains certain keyword. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "71aeb10e5fda9040",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:13.813566Z",
     "start_time": "2024-01-25T14:27:13.693010Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.runnables import Runnable\n",
    "\n",
    "\n",
    "class CheckKeywordsRunnable(Runnable):\n",
    "    def invoke(self, input, config=None, **kwargs):\n",
    "        text = input[\"text\"]\n",
    "        keywords = input[\"keywords\"].split(\",\")\n",
    "\n",
    "        for keyword in keywords:\n",
    "            if keyword.strip() in text:\n",
    "                return True\n",
    "\n",
    "        return False\n",
    "\n",
    "\n",
    "print(CheckKeywordsRunnable().invoke({\"text\": \"This is a proprietary message\", \"keywords\": \"proprietary\"}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a0725d977f5589b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Guardrails Configuration \n",
    "\n",
    "Now, let's create a guardrails configuration that uses the `CheckKeywords` runnable as part of an input rail flow. To achieve this, you need to register an instance of `CheckKeywords` as an action. In the snippets below, we register it as the `check_keywords` action. We can then use this action inside the `check proprietary keywords` flow, which is used as an input rail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a27c15cf3919fa5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:13.820255Z",
     "start_time": "2024-01-25T14:27:13.814191Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing config/rails.co\n"
     ]
    }
   ],
   "source": [
    "%%writefile config/rails.co\n",
    "\n",
    "define flow check proprietary keywords\n",
    "  $keywords = \"proprietary\"\n",
    "  $has_keywords = execute check_keywords(text=$user_message, keywords=$keywords)\n",
    "  \n",
    "  if $has_keywords\n",
    "    bot refuse to respond\n",
    "    stop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "53403afb1e1a4b9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:13.821992Z",
     "start_time": "2024-01-25T14:27:13.817004Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing config/config.yml\n"
     ]
    }
   ],
   "source": [
    "%%writefile config/config.yml\n",
    "models:\n",
    " - type: main\n",
    "   engine: openai\n",
    "   model: gpt-3.5-turbo-instruct\n",
    "\n",
    "rails:\n",
    "  input:\n",
    "    flows:\n",
    "      - check proprietary keywords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2adca21d94e54b9",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from nemoguardrails import LLMRails, RailsConfig\n",
    "\n",
    "config = RailsConfig.from_path(\"./config\")\n",
    "rails = LLMRails(config)\n",
    "\n",
    "rails.register_action(CheckKeywordsRunnable(), \"check_keywords\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ade12682dd9d8f0e",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Testing\n",
    "\n",
    "Let's give this a try. If we invoke the guardrails configuration with a message that contains the \"proprietary\" keyword, the returned response is \"I'm sorry, I can't respond to that\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "394311174e678d96",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:27:18.524958Z",
     "start_time": "2024-01-25T14:27:18.518176Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm sorry, I can't respond to that.\n"
     ]
    }
   ],
   "source": [
    "response = rails.generate(\"Give me some proprietary information.\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b457ce6e2957fd",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "On the other hand, a message which does not hit the input rail, will proceed as usual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "70409a3aafe89e95",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-25T14:29:15.370273Z",
     "start_time": "2024-01-25T14:29:14.322661Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The result for 2+2 is 4. This is a basic addition problem that can also be written as 2 plus 2 equals 4, or two plus two equals four. The answer is a basic fact that is often taught in early elementary school and is an important building block for more complex mathematical concepts.\n"
     ]
    }
   ],
   "source": [
    "response = rails.generate(\"What is the result for 2+2?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39bd84e0a3fb94e1",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Conclusion\n",
    "\n",
    "In this guide, you learned how to register a custom `Runnable` as an action and use it inside a guardrails configuration. This guide uses a basic implementation of a `Runnable`. However, you can register any type of `Runnable`, including ones that make calls to the LLM, 3rd party APIs or vector stores."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
